// 此代码由DeepSeek生成的代码修改而来
#pragma once

#include <vector>
#include <stdint.h>
#include <random>
#include <stdexcept>
#include <algorithm>

class Iterator
{
public:
    explicit Iterator(int64_t data_size);
    std::vector<std::vector<double>> get_rand_batch(const std::vector<std::vector<double>> &data, int64_t batch);
    template<typename feature_type, typename label_type>
    void get_batch(const std::vector<std::vector<feature_type>> &feature, const std::vector<std::vector<label_type>> &label, 
        std::vector<std::vector<feature_type>> &feature_batch, std::vector<std::vector<label_type>> &label_batch, int64_t batch_size)
    {
        // 1. 参数检查
        if (feature.empty() || label.empty() || feature.size() != label.size()) {
            throw std::invalid_argument("Feature and label sizes must match");
        }
        if (batch_size <= 0) {
            throw std::invalid_argument("Batch size must be positive");
        }

        // 2. 惰性初始化索引
        if (indices_.size() != feature.size()) {
            reset(feature.size());
        }

        // 3. 检查是否需要重新洗牌
        if (cursor_ + batch_size > feature.size()) {
            std::shuffle(indices_.begin(), indices_.end(), gen_);
            cursor_ = 0;
        }

        // 4. 预分配批次内存
        feature_batch.clear();
        label_batch.clear();
        feature_batch.reserve(batch_size);
        label_batch.reserve(batch_size);

        // 5. 填充批次数据
        for (; cursor_ < indices_.size() && feature_batch.size() < static_cast<size_t>(batch_size); ++cursor_) {
            feature_batch.push_back(feature[indices_[cursor_]]);
            label_batch.push_back(label[indices_[cursor_]]);
        }
    }
private:
    void reset(size_t data_size);
private:
    std::vector<size_t> indices_;  // 全局随机索引
    size_t cursor_ = 0;            // 当前读取位置
    std::mt19937 gen_;             // 复用随机引擎
};
